/**
 * @since 1.0.0
 */
import * as AiError from "@effect/ai/AiError"
import * as IdGenerator from "@effect/ai/IdGenerator"
import * as LanguageModel from "@effect/ai/LanguageModel"
import * as AiModel from "@effect/ai/Model"
import type * as Prompt from "@effect/ai/Prompt"
import type * as Response from "@effect/ai/Response"
import { addGenAIAnnotations } from "@effect/ai/Telemetry"
import type * as Tokenizer from "@effect/ai/Tokenizer"
import * as Tool from "@effect/ai/Tool"
import * as Arr from "effect/Array"
import * as Context from "effect/Context"
import * as DateTime from "effect/DateTime"
import * as Effect from "effect/Effect"
import * as Encoding from "effect/Encoding"
import { dual } from "effect/Function"
import * as Layer from "effect/Layer"
import * as Predicate from "effect/Predicate"
import * as Stream from "effect/Stream"
import type { Span } from "effect/Tracer"
import type { Mutable, Simplify } from "effect/Types"
import { AnthropicClient, type MessageStreamEvent } from "./AnthropicClient.js"
import * as AnthropicTokenizer from "./AnthropicTokenizer.js"
import * as AnthropicTool from "./AnthropicTool.js"
import type * as Generated from "./Generated.js"
import * as InternalUtilities from "./internal/utilities.js"

/**
 * @since 1.0.0
 * @category Models
 */
export type Model = typeof Generated.Model.Encoded

// =============================================================================
// Configuration
// =============================================================================

/**
 * @since 1.0.0
 * @category Context
 */
export class Config extends Context.Tag("@effect/ai-anthropic/AnthropicLanguageModel/Config")<
  Config,
  Config.Service
>() {
  /**
   * @since 1.0.0
   */
  static readonly getOrUndefined: Effect.Effect<typeof Config.Service | undefined> = Effect.map(
    Effect.context<never>(),
    (context) => context.unsafeMap.get(Config.key)
  )
}

/**
 * @since 1.0.0
 */
export declare namespace Config {
  /**
   * @since 1.0.0
   * @category Configuration
   */
  export interface Service extends
    Simplify<
      Partial<
        Omit<
          typeof Generated.CreateMessageParams.Encoded,
          "messages" | "tools" | "tool_choice" | "stream"
        >
      >
    >
  {
    readonly disableParallelToolCalls?: boolean
  }
}

// =============================================================================
// Anthropic Provider Options / Metadata
// =============================================================================

/**
 * @since 1.0.0
 * @category Provider Metadata
 */
export type AnthropicReasoningInfo = {
  readonly type: "thinking"
  /**
   * Thinking content as an encrypted string, which is used to verify
   * that thinking content was indeed generated by Anthropic's API.
   */
  readonly signature: typeof Generated.ResponseThinkingBlock.fields.thinking.Encoded
} | {
  readonly type: "redacted_thinking"
  /**
   * Thinking content which was flagged by Anthropic's safety systems, and
   * was therefore encrypted.
   */
  readonly redactedData: typeof Generated.RequestRedactedThinkingBlock.fields.data.Encoded
}

/**
 * @since 1.0.0
 * @category Provider Options
 */
declare module "@effect/ai/Prompt" {
  export interface SystemMessageOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
    } | undefined
  }

  export interface UserMessageOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
    } | undefined
  }

  export interface AssistantMessageOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
    } | undefined
  }

  export interface ToolMessageOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
    } | undefined
  }

  export interface TextPartOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
    } | undefined
  }

  export interface ReasoningPartOptions extends ProviderOptions {
    readonly anthropic?:
      | Simplify<
        AnthropicReasoningInfo & {
          /**
           * A breakpoint which marks the end of reusable content eligible for caching.
           */
          readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
        }
      >
      | undefined
  }

  export interface FilePartOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
      /**
       * Whether or not citations should be enabled for the file part.
       */
      readonly citations?: typeof Generated.RequestCitationsConfig.Encoded | undefined
      /**
       * A custom title to provide to the document. If omitted, the file part's
       * `fileName` property will be used.
       */
      readonly documentTitle?: string | undefined
      /**
       * Additional context about the document that will be forwarded to the
       * large language model, but will not be used towards cited content.
       *
       * Useful for storing additional document metadata as text or stringified JSON.
       */
      readonly documentContext?: string | undefined
    } | undefined
  }

  export interface ToolCallPartOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
    } | undefined
  }

  export interface ToolResultPartOptions extends ProviderOptions {
    readonly anthropic?: {
      /**
       * A breakpoint which marks the end of reusable content eligible for caching.
       */
      readonly cacheControl?: typeof Generated.CacheControlEphemeral.Encoded | undefined
    } | undefined
  }
}

declare module "@effect/ai/Response" {
  export interface ReasoningPartMetadata extends ProviderMetadata {
    readonly anthropic?: AnthropicReasoningInfo | undefined
  }

  export interface ReasoningStartPartMetadata extends ProviderMetadata {
    readonly anthropic?: AnthropicReasoningInfo | undefined
  }

  export interface ReasoningDeltaPartMetadata extends ProviderMetadata {
    readonly anthropic?: AnthropicReasoningInfo | undefined
  }

  export interface FinishPartMetadata extends ProviderMetadata {
    readonly anthropic?: {
      /**
       * Additional usage information provided by the Anthropic API.
       */
      readonly usage?: Generated.BetaUsage | undefined
      /**
       * Which custom stop sequence was generated, if any.
       *
       * If one of the custom user-defined stop sequences was generated, the
       * value will be a `string` with that stop sequence.
       */
      readonly stopSequence?: string | undefined
    } | undefined
  }

  export interface DocumentSourcePartMetadata extends ProviderMetadata {
    readonly anthropic?: {
      readonly source: "document"
      readonly type: "char_location"
      /**
       * The text that was cited in the response.
       */
      readonly citedText: string
      /**
       * The 0-indexed starting position of the characters that were cited.
       */
      readonly startCharIndex: number
      /**
       * The exclusive ending position of the characters that were cited.
       */
      readonly endCharIndex: number
    } | {
      readonly source: "document"
      readonly type: "page_location"
      /**
       * The text that was cited in the response.
       */
      readonly citedText: string
      /**
       * The 1-indexed starting page of pages that were cited.
       */
      readonly startPageNumber: number
      /**
       * The exclusive ending position of the pages that were cited.
       */
      readonly endPageNumber: number
    } | undefined
  }

  export interface UrlSourcePartMetadata extends ProviderMetadata {
    readonly anthropic?: {
      readonly source: "url"
      /**
       * Up to 150 characters of the text content that was referenced from the
       * URL source material.
       */
      readonly citedText: string
      /**
       * An internal reference that must be passed back to the Anthropic API
       * during multi-turn conversations.
       */
      readonly encryptedIndex: string
    } | undefined
  }
}

// =============================================================================
// Anthropic Language Model
// =============================================================================

/**
 * @since 1.0.0
 * @category Ai Models
 */
export const model = (
  model: (string & {}) | Model,
  config?: Omit<Config.Service, "model">
): AiModel.Model<"anthropic", LanguageModel.LanguageModel, AnthropicClient> =>
  AiModel.make("anthropic", layer({ model, config }))

/**
 * @since 1.0.0
 * @category Ai Models
 */
export const modelWithTokenizer = (
  model: (string & {}) | Model,
  config?: Omit<Config.Service, "model">
): AiModel.Model<"anthropic", LanguageModel.LanguageModel | Tokenizer.Tokenizer, AnthropicClient> =>
  AiModel.make("anthropic", layerWithTokenizer({ model, config }))

/**
 * @since 1.0.0
 * @category Constructors
 */
export const make = Effect.fnUntraced(function*(options: {
  readonly model: (string & {}) | Model
  readonly config?: Omit<Config.Service, "model">
}) {
  const client = yield* AnthropicClient

  const makeRequest = Effect.fnUntraced(
    function*(providerOptions: LanguageModel.ProviderOptions) {
      const context = yield* Effect.context<never>()
      const config = { model: options.model, ...options.config, ...context.unsafeMap.get(Config.key) }
      const { betas: messageBetas, messages, system } = yield* prepareMessages(providerOptions)
      const { betas: toolBetas, toolChoice, tools } = yield* prepareTools(providerOptions, config)
      const responseFormat = providerOptions.responseFormat
      const request: typeof Generated.BetaCreateMessageParams.Encoded = {
        max_tokens: 4096,
        ...config,
        system,
        messages,
        tools: responseFormat.type === "text"
          ? tools
          : [{
            name: responseFormat.objectName,
            description: Tool.getDescriptionFromSchemaAst(responseFormat.schema.ast) ?? "Respond with a JSON object",
            input_schema: Tool.getJsonSchemaFromSchemaAst(responseFormat.schema.ast) as any
          }],
        tool_choice: responseFormat.type === "text"
          ? toolChoice
          : {
            type: "tool",
            name: responseFormat.objectName,
            disable_parallel_tool_use: true
          }
      }
      return { betas: new Set([...messageBetas, ...toolBetas]), request }
    }
  )

  return yield* LanguageModel.make({
    generateText: Effect.fnUntraced(
      function*(options) {
        const { betas, request } = yield* makeRequest(options)
        annotateRequest(options.span, request)
        const anthropicBeta = betas.size > 0 ? Array.from(betas).join(",") : undefined
        const rawResponse = yield* client.createMessage({
          params: { "anthropic-beta": anthropicBeta },
          payload: request
        })
        annotateResponse(options.span, rawResponse)
        return yield* makeResponse(rawResponse, options)
      }
    ),
    streamText: Effect.fnUntraced(
      function*(options) {
        const { betas, request } = yield* makeRequest(options)
        annotateRequest(options.span, request)
        const anthropicBeta = betas.size > 0 ? Array.from(betas).join(",") : undefined
        return client.createMessageStream({
          params: { "anthropic-beta": anthropicBeta },
          payload: request
        })
      },
      (effect, options) =>
        effect.pipe(
          Effect.flatMap((stream) => makeStreamResponse(stream, options)),
          Stream.unwrap,
          Stream.map((response) => {
            annotateStreamResponse(options.span, response)
            return response
          })
        )
    )
  })
})

/**
 * @since 1.0.0
 * @category Layers
 */
export const layer = (options: {
  readonly model: (string & {}) | Model
  readonly config?: Omit<Config.Service, "model">
}): Layer.Layer<LanguageModel.LanguageModel, never, AnthropicClient> =>
  Layer.effect(LanguageModel.LanguageModel, make({ model: options.model, config: options.config }))

/**
 * @since 1.0.0
 * @category Layers
 */
export const layerWithTokenizer = (options: {
  readonly model: (string & {}) | Model
  readonly config?: Omit<Config.Service, "model">
}): Layer.Layer<LanguageModel.LanguageModel | Tokenizer.Tokenizer, never, AnthropicClient> =>
  Layer.merge(layer(options), AnthropicTokenizer.layer)

/**
 * @since 1.0.0
 * @category Configuration
 */
export const withConfigOverride: {
  (config: Config.Service): <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>
  <A, E, R>(self: Effect.Effect<A, E, R>, config: Config.Service): Effect.Effect<A, E, R>
} = dual<
  (config: Config.Service) => <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, R>,
  <A, E, R>(self: Effect.Effect<A, E, R>, config: Config.Service) => Effect.Effect<A, E, R>
>(2, (self, overrides) =>
  Effect.flatMap(
    Config.getOrUndefined,
    (config) => Effect.provideService(self, Config, { ...config, ...overrides })
  ))

// =============================================================================
// Prompt Conversion
// =============================================================================

const prepareMessages: (options: LanguageModel.ProviderOptions) => Effect.Effect<{
  readonly betas: ReadonlySet<string>
  readonly system: ReadonlyArray<typeof Generated.BetaRequestTextBlock.Encoded> | undefined
  readonly messages: ReadonlyArray<typeof Generated.BetaInputMessage.Encoded>
}, AiError.AiError> = Effect.fnUntraced(function*(options) {
  const betas = new Set<string>()
  const groups = groupMessages(options.prompt)

  let system: Array<typeof Generated.BetaRequestTextBlock.Encoded> | undefined = undefined
  const messages: Array<typeof Generated.BetaInputMessage.Encoded> = []

  for (let i = 0; i < groups.length; i++) {
    const group = groups[i]
    const isLastGroup = i === groups.length - 1

    switch (group.type) {
      case "system": {
        system = group.messages.map((message) => ({
          type: "text",
          text: message.content,
          cache_control: getCacheControl(message)
        }))
        break
      }

      case "user": {
        const content: Array<typeof Generated.BetaInputContentBlock.Encoded> = []

        for (const message of group.messages) {
          switch (message.role) {
            case "user": {
              for (let j = 0; j < message.content.length; j++) {
                const part = message.content[j]
                const isLastPart = j === message.content.length - 1

                // Attempt to get the cache control from the part first. If
                // the part does not have cache control defined and we are
                // evaluating the last part for this message, also check the
                // message for cache control.
                const cacheControl = getCacheControl(part) ?? (
                  isLastPart ? getCacheControl(message) : undefined
                )

                switch (part.type) {
                  case "text": {
                    content.push({
                      type: "text",
                      text: part.text,
                      cache_control: cacheControl
                    })
                    break
                  }

                  case "file": {
                    if (part.mediaType.startsWith("image/")) {
                      const source = part.data instanceof URL ?
                        {
                          type: "url",
                          url: part.data.toString()
                        } as const :
                        {
                          type: "base64",
                          media_type: part.mediaType === "image/*"
                            ? "image/jpeg"
                            : (part.mediaType as typeof Generated.Base64ImageSourceMediaType.Encoded),
                          data: typeof part.data === "string" ? part.data : Encoding.encodeBase64(part.data)
                        } as const

                      content.push({
                        type: "image",
                        source,
                        cache_control: cacheControl
                      })
                    } else if (part.mediaType === "application/pdf" || part.mediaType === "text/plain") {
                      if (part.mediaType === "application/pdf") {
                        betas.add("pdfs-2024-09-25")
                      }

                      const enableCitations = shouldEnableCitations(part)
                      const documentOptions = getDocumentMetadata(part)

                      const source = part.data instanceof URL
                        ? {
                          type: "url",
                          url: part.data.toString()
                        } as const
                        : part.mediaType === "application/pdf"
                        ? {
                          type: "base64",
                          media_type: "application/pdf",
                          data: typeof part.data === "string"
                            ? part.data
                            : Encoding.encodeBase64(part.data)
                        } as const
                        : {
                          type: "text",
                          media_type: "text/plain",
                          data: typeof part.data === "string"
                            ? part.data
                            : Encoding.encodeBase64(part.data)
                        } as const

                      content.push({
                        type: "document",
                        source,
                        title: documentOptions?.title ?? part.fileName,
                        ...(documentOptions?.context ? { context: documentOptions.context } : undefined),
                        ...(enableCitations ? { citations: { enabled: true } } : undefined),
                        cache_control: cacheControl
                      })
                    } else {
                      return yield* new AiError.MalformedInput({
                        module: "AnthropicLanguageModel",
                        method: "prepareMessages",
                        description: `Detected unsupported media type for file: '${part.mediaType}'`
                      })
                    }
                    break
                  }
                }
              }

              break
            }

            // TODO: advanced tool result content parts
            case "tool": {
              for (let j = 0; j < message.content.length; j++) {
                const part = message.content[j]
                const isLastPart = j === message.content.length - 1

                // Attempt to get the cache control from the part first. If
                // the part does not have cache control defined and we are
                // evaluating the last part for this message, also check the
                // message for cache control.
                const cacheControl = getCacheControl(part) ?? (
                  isLastPart ? getCacheControl(message) : undefined
                )

                content.push({
                  type: "tool_result",
                  tool_use_id: part.id,
                  content: JSON.stringify(part.result),
                  is_error: part.isFailure,
                  cache_control: cacheControl
                })
              }

              break
            }
          }
        }

        messages.push({ role: "user", content })

        break
      }

      case "assistant": {
        const content: Array<typeof Generated.BetaInputContentBlock.Encoded> = []

        for (let j = 0; j < group.messages.length; j++) {
          const message = group.messages[j]
          const isLastMessage = j === group.messages.length - 1

          for (let k = 0; k < message.content.length; k++) {
            const part = message.content[k]
            const isLastPart = k === message.content.length - 1

            // Attempt to get the cache control from the part first. If
            // the part does not have cache control defined and we are
            // evaluating the last part for this message, also check the
            // message for cache control.
            const cacheControl = getCacheControl(part) ?? (
              isLastPart ? getCacheControl(message) : undefined
            )

            switch (part.type) {
              case "text": {
                content.push({
                  type: "text",
                  // Anthropic does not allow trailing whitespace in assistant
                  // content blocks
                  text: isLastGroup && isLastMessage && isLastPart
                    ? part.text.trim()
                    : part.text
                })
                break
              }

              case "reasoning": {
                const options = part.options.anthropic
                if (Predicate.isNotUndefined(options)) {
                  if (options.type === "thinking") {
                    content.push({
                      type: "thinking",
                      thinking: part.text,
                      signature: options.signature
                    })
                  } else {
                    content.push({
                      type: "redacted_thinking",
                      data: options.redactedData
                    })
                  }
                }
                break
              }

              case "tool-call": {
                if (part.providerExecuted) {
                  if (part.name === "AnthropicCodeExecution") {
                    content.push({
                      type: "server_tool_use",
                      id: part.id,
                      name: "code_execution",
                      input: part.params as any,
                      cache_control: cacheControl
                    })
                  }
                  if (part.name === "AnthropicWebSearch") {
                    content.push({
                      type: "server_tool_use",
                      id: part.id,
                      name: "web_search",
                      input: part.params as any,
                      cache_control: cacheControl
                    })
                  }
                } else {
                  content.push({
                    type: "tool_use",
                    id: part.id,
                    name: part.name,
                    input: part.params as any,
                    cache_control: cacheControl
                  })
                }
                break
              }

              case "tool-result": {
                if (part.name === "AnthropicCodeExecution") {
                  content.push({
                    type: "code_execution_tool_result",
                    tool_use_id: part.id,
                    content: part.result as any,
                    cache_control: cacheControl
                  })
                } else if (part.name === "AnthropicWebSearch") {
                  content.push({
                    type: "web_search_tool_result",
                    tool_use_id: part.id,
                    content: part.result as any,
                    cache_control: cacheControl
                  })
                } else {
                  return yield* new AiError.MalformedInput({
                    module: "AnthropicLanguageModel",
                    method: "prepareMessages",
                    description: `Provider executed tool result for tool ${part.name} is not supported in prompt`
                  })
                }
              }
            }
          }
        }

        messages.push({ role: "assistant", content })

        break
      }
    }
  }

  return {
    system,
    messages,
    betas
  }
})

// =============================================================================
// Response Conversion
// =============================================================================

const makeResponse: (
  response: Generated.BetaMessage,
  options: LanguageModel.ProviderOptions
) => Effect.Effect<
  Array<Response.PartEncoded>,
  never,
  IdGenerator.IdGenerator
> = Effect.fnUntraced(
  function*(response, options) {
    const idGenerator = yield* IdGenerator.IdGenerator
    const parts: Array<Response.PartEncoded> = []
    const citableDocuments = extractCitableDocuments(options.prompt)

    parts.push({
      type: "response-metadata",
      id: response.id,
      modelId: response.model,
      timestamp: DateTime.formatIso(yield* DateTime.now)
    })

    for (const part of response.content) {
      switch (part.type) {
        case "text": {
          // The text parts should only be added to the response here if the
          // response format is `"text"`. If the response format is `"json"`,
          // then the text parts must instead be added to the response when a
          // tool call is received.
          if (options.responseFormat.type === "text") {
            parts.push({
              type: "text",
              text: part.text
            })

            if (Predicate.isNotNullable(part.citations)) {
              for (const citation of part.citations) {
                const source = yield* processCitation(citation, citableDocuments, idGenerator)
                if (Predicate.isNotUndefined(source)) {
                  parts.push(source)
                }
              }
            }
          }

          break
        }

        case "thinking": {
          parts.push({
            type: "reasoning",
            text: part.thinking,
            metadata: { anthropic: { type: "thinking", signature: part.signature } }
          })
          break
        }

        case "redacted_thinking": {
          parts.push({
            type: "reasoning",
            text: "",
            metadata: { anthropic: { type: "redacted_thinking", redactedData: part.data } }
          })
          break
        }

        case "tool_use": {
          // When a `"json"` response format is requested, the JSON that we need
          // will be returned by the tool call injected into the request
          if (options.responseFormat.type === "json") {
            parts.push({
              type: "text",
              text: JSON.stringify(part.input)
            })
          } else {
            const providerTool = AnthropicTool.getProviderDefinedToolName(part.name)
            const name = Predicate.isNotUndefined(providerTool) ? providerTool : part.name
            const providerName = Predicate.isNotUndefined(providerTool) ? part.name : undefined
            parts.push({
              type: "tool-call",
              id: part.id,
              name,
              params: part.input,
              providerName,
              providerExecuted: false
            })
          }

          break
        }

        case "server_tool_use": {
          const providerTool = AnthropicTool.getProviderDefinedToolName(part.name)
          if (Predicate.isNotUndefined(providerTool)) {
            parts.push({
              type: "tool-call",
              id: part.id,
              name: providerTool,
              params: part.input,
              providerName: part.name,
              providerExecuted: true
            })
          }

          break
        }

        case "bash_code_execution_tool_result": {
          const isFailure = part.content.type === "bash_code_execution_tool_result_error"
          parts.push({
            type: "tool-result",
            id: part.tool_use_id,
            name: "AnthropicCodeExecution",
            isFailure,
            result: part.content,
            providerName: "code_execution",
            providerExecuted: true
          })
          break
        }

        case "code_execution_tool_result": {
          const isFailure = part.content.type === "code_execution_tool_result_error"
          parts.push({
            type: "tool-result",
            id: part.tool_use_id,
            name: "AnthropicCodeExecution",
            isFailure,
            result: part.content,
            providerName: "code_execution",
            providerExecuted: true
          })
          break
        }

        case "text_editor_code_execution_tool_result": {
          const isFailure = part.content.type === "text_editor_code_execution_tool_result_error"
          parts.push({
            type: "tool-result",
            id: part.tool_use_id,
            name: "AnthropicCodeExecution",
            isFailure,
            result: part.content,
            providerName: "code_execution",
            providerExecuted: true
          })
          break
        }

        case "web_search_tool_result": {
          const isFailure = !Array.isArray(part.content)
          parts.push({
            type: "tool-result",
            id: part.tool_use_id,
            name: "AnthropicWebSearch",
            isFailure,
            result: part.content,
            providerName: "web_search",
            providerExecuted: true
          })
          break
        }
      }
    }

    // Anthropic always returns a non-null `stop_reason` for non-streaming responses
    const finishReason = InternalUtilities.resolveFinishReason(
      response.stop_reason!,
      options.responseFormat.type === "json"
    )

    parts.push({
      type: "finish",
      reason: finishReason,
      usage: {
        inputTokens: response.usage.input_tokens,
        outputTokens: response.usage.output_tokens,
        totalTokens: response.usage.input_tokens + response.usage.output_tokens,
        cachedInputTokens: response.usage.cache_read_input_tokens ?? undefined
      },
      metadata: {
        anthropic: {
          usage: response.usage,
          stopSequence: response.stop_sequence ?? undefined
        }
      }
    })

    return parts
  }
)

const makeStreamResponse: (
  stream: Stream.Stream<MessageStreamEvent, AiError.AiError>,
  options: LanguageModel.ProviderOptions
) => Effect.Effect<
  Stream.Stream<Response.StreamPartEncoded, AiError.AiError>,
  never,
  IdGenerator.IdGenerator
> = Effect.fnUntraced(
  function*(stream, options) {
    const idGenerator = yield* IdGenerator.IdGenerator
    const citableDocuments = extractCitableDocuments(options.prompt)

    // Setup all requisite state for the streaming response
    let finishReason: Response.FinishReason = "unknown"
    const contentBlocks: Record<
      number,
      | {
        readonly type: "text"
      }
      | {
        readonly type: "reasoning"
      }
      | {
        readonly type: "tool-call"
        readonly id: string
        readonly name: string
        params: string
        readonly providerName: string | undefined
        readonly providerExecuted: boolean
      }
    > = {}
    let blockType:
      | "text"
      | "thinking"
      | "redacted_thinking"
      | "tool_use"
      | "server_tool_use"
      | "web_fetch_tool_result"
      | "web_search_tool_result"
      | "code_execution_tool_result"
      | "bash_code_execution_tool_result"
      | "text_editor_code_execution_tool_result"
      | "mcp_tool_use"
      | "mcp_tool_result"
      | "container_upload"
      | undefined = undefined
    const usage: Mutable<typeof Response.Usage.Encoded> = {
      inputTokens: undefined,
      outputTokens: undefined,
      totalTokens: undefined
    }
    let metaUsage: Generated.BetaUsage | undefined = undefined
    let stopSequence: string | undefined = undefined

    return stream.pipe(
      Stream.mapEffect(Effect.fnUntraced(function*(event) {
        const parts: Array<Response.StreamPartEncoded> = []

        switch (event.type) {
          case "ping": {
            break
          }

          case "message_start": {
            // Track usage metadata
            usage.inputTokens = event.message.usage.input_tokens
            metaUsage = event.message.usage

            // Track response metadata
            parts.push({
              type: "response-metadata",
              id: event.message.id,
              modelId: event.message.model,
              timestamp: DateTime.formatIso(yield* DateTime.now)
            })

            break
          }

          case "message_delta": {
            // Track usage metadata
            if (Predicate.isNotNullable(event.usage.output_tokens)) {
              usage.outputTokens = event.usage.output_tokens
            }
            usage.totalTokens = (usage.inputTokens ?? 0) + (event.usage.output_tokens ?? 0)

            // Track stop sequence metadata
            if (Predicate.isNotNullable(event.delta.stop_sequence)) {
              stopSequence = event.delta.stop_sequence
            }

            // Track the response finish reason
            if (Predicate.isNotNullable(event.delta.stop_reason)) {
              finishReason = InternalUtilities.resolveFinishReason(event.delta.stop_reason)
            }

            break
          }

          case "message_stop": {
            parts.push({
              type: "finish",
              reason: finishReason,
              usage,
              metadata: { anthropic: { usage: metaUsage, stopSequence } }
            })

            break
          }

          case "content_block_start": {
            blockType = event.content_block.type

            switch (event.content_block.type) {
              case "text": {
                contentBlocks[event.index] = { type: "text" }

                parts.push({
                  type: "text-start",
                  id: event.index.toString()
                })

                break
              }

              case "thinking": {
                contentBlocks[event.index] = { type: "reasoning" }

                parts.push({
                  type: "reasoning-start",
                  id: event.index.toString()
                })

                break
              }

              case "redacted_thinking": {
                contentBlocks[event.index] = { type: "reasoning" }

                parts.push({
                  type: "reasoning-start",
                  id: event.index.toString(),
                  metadata: {
                    anthropic: {
                      type: "redacted_thinking",
                      redactedData: event.content_block.data
                    }
                  }
                })

                break
              }

              case "tool_use": {
                const toolName = event.content_block.name
                const providerTool = AnthropicTool.getProviderDefinedToolName(toolName)
                const name = Predicate.isNotUndefined(providerTool) ? providerTool : toolName
                const providerName = Predicate.isNotUndefined(providerTool) ? toolName : undefined

                contentBlocks[event.index] = {
                  type: "tool-call",
                  id: event.content_block.id,
                  name,
                  params: "",
                  providerName,
                  providerExecuted: false
                }

                parts.push({
                  type: "tool-params-start",
                  id: event.content_block.id,
                  name: toolName,
                  providerName,
                  providerExecuted: false
                })

                break
              }

              case "server_tool_use": {
                const toolName = event.content_block.name
                const providerTool = AnthropicTool.getProviderDefinedToolName(toolName)
                if (Predicate.isNotUndefined(providerTool)) {
                  contentBlocks[event.index] = {
                    type: "tool-call",
                    id: event.content_block.id,
                    name: providerTool,
                    params: "",
                    providerName: toolName,
                    providerExecuted: true
                  }

                  parts.push({
                    type: "tool-params-start",
                    id: event.content_block.id,
                    name: providerTool,
                    providerName: toolName,
                    providerExecuted: true
                  })
                }

                break
              }

              case "bash_code_execution_tool_result": {
                const toolUseId = event.content_block.tool_use_id
                const content = event.content_block.content
                const isFailure = content.type === "bash_code_execution_tool_result_error"
                parts.push({
                  type: "tool-result",
                  id: toolUseId,
                  name: "AnthropicCodeExecution",
                  isFailure,
                  result: content,
                  providerName: "code_execution",
                  providerExecuted: true
                })
                break
              }

              case "code_execution_tool_result": {
                const toolUseId = event.content_block.tool_use_id
                const content = event.content_block.content
                const isFailure = content.type === "code_execution_tool_result_error"
                parts.push({
                  type: "tool-result",
                  id: toolUseId,
                  name: "AnthropicCodeExecution",
                  isFailure,
                  result: content,
                  providerName: "code_execution",
                  providerExecuted: true
                })
                break
              }

              case "text_editor_code_execution_tool_result": {
                const toolUseId = event.content_block.tool_use_id
                const content = event.content_block.content
                const isFailure = content.type === "text_editor_code_execution_tool_result_error"
                parts.push({
                  type: "tool-result",
                  id: toolUseId,
                  name: "AnthropicCodeExecution",
                  isFailure,
                  result: content,
                  providerName: "code_execution",
                  providerExecuted: true
                })
                break
              }

              case "web_search_tool_result": {
                const toolUseId = event.content_block.tool_use_id
                const content = event.content_block.content
                const isFailure = !Array.isArray(content)
                parts.push({
                  type: "tool-result",
                  id: toolUseId,
                  name: "AnthropicWebSearch",
                  isFailure,
                  result: content,
                  providerName: "web_search",
                  providerExecuted: true
                })
                break
              }
            }

            break
          }

          case "content_block_delta": {
            switch (event.delta.type) {
              case "text_delta": {
                parts.push({
                  type: "text-delta",
                  id: event.index.toString(),
                  delta: event.delta.text
                })

                break
              }

              case "thinking_delta": {
                parts.push({
                  type: "reasoning-delta",
                  id: event.index.toString(),
                  delta: event.delta.thinking
                })

                break
              }

              case "signature_delta": {
                if (blockType === "thinking") {
                  parts.push({
                    type: "reasoning-delta",
                    id: event.index.toString(),
                    delta: "",
                    metadata: {
                      anthropic: {
                        type: "thinking",
                        signature: event.delta.signature
                      }
                    }
                  })
                }

                break
              }

              case "input_json_delta": {
                const contentBlock = contentBlocks[event.index]
                const delta = event.delta.partial_json

                if (contentBlock.type === "tool-call") {
                  parts.push({
                    type: "tool-params-delta",
                    id: contentBlock.id,
                    delta
                  })

                  contentBlock.params += delta
                }

                break
              }

              case "citations_delta": {
                const citation = event.delta.citation

                const source = yield* processCitation(citation, citableDocuments, idGenerator)
                if (Predicate.isNotUndefined(source)) {
                  parts.push(source)
                }
              }
            }

            break
          }

          case "content_block_stop": {
            if (Predicate.isNotNullable(contentBlocks[event.index])) {
              const contentBlock = contentBlocks[event.index]

              switch (contentBlock.type) {
                case "text": {
                  parts.push({
                    type: "text-end",
                    id: event.index.toString()
                  })
                  break
                }

                case "reasoning": {
                  parts.push({
                    type: "reasoning-end",
                    id: event.index.toString()
                  })
                  break
                }

                case "tool-call": {
                  parts.push({
                    type: "tool-params-end",
                    id: contentBlock.id
                  })

                  const toolName = contentBlock.name
                  // If the tool call has no parameters, an empty string is returned
                  const toolParams = contentBlock.params.length === 0 ? "{}" : contentBlock.params

                  const parsedParams = yield* Effect.try({
                    try: () => Tool.unsafeSecureJsonParse(toolParams),
                    catch: (cause) =>
                      new AiError.MalformedOutput({
                        module: "AnthropicLanguageModel",
                        method: "makeStreamResponse",
                        description: "Failed to securely parse tool call parameters " +
                          `for tool '${toolName}':\nParameters: ${toolParams}`,
                        cause
                      })
                  })

                  parts.push({
                    type: "tool-call",
                    id: contentBlock.id,
                    name: toolName,
                    params: parsedParams,
                    providerName: contentBlock.providerName,
                    providerExecuted: contentBlock.providerExecuted
                  })

                  break
                }
              }

              delete contentBlocks[event.index]
            }

            blockType = undefined

            break
          }

          case "error": {
            parts.push({ type: "error", error: event.error })

            break
          }
        }

        return parts
      })),
      Stream.flattenIterables
    )
  }
)

// =============================================================================
// Telemetry
// =============================================================================

const annotateRequest = (
  span: Span,
  request: typeof Generated.BetaCreateMessageParams.Encoded
): void => {
  addGenAIAnnotations(span, {
    system: "anthropic",
    operation: { name: "chat" },
    request: {
      model: request.model,
      temperature: request.temperature,
      topK: request.top_k,
      topP: request.top_p,
      maxTokens: request.max_tokens,
      stopSequences: Arr.ensure(request.stop_sequences).filter(
        Predicate.isNotNullable
      )
    }
  })
}

const annotateResponse = (span: Span, response: Generated.BetaMessage): void => {
  addGenAIAnnotations(span, {
    response: {
      id: response.id,
      model: response.model,
      finishReasons: response.stop_reason ? [response.stop_reason] : undefined
    },
    usage: {
      inputTokens: response.usage.input_tokens,
      outputTokens: response.usage.output_tokens
    }
  })
}

const annotateStreamResponse = (span: Span, part: Response.StreamPartEncoded) => {
  if (part.type === "response-metadata") {
    addGenAIAnnotations(span, {
      response: {
        id: part.id,
        model: part.modelId
      }
    })
  }
  if (part.type === "finish") {
    addGenAIAnnotations(span, {
      response: {
        finishReasons: [part.reason]
      },
      usage: {
        inputTokens: part.usage.inputTokens,
        outputTokens: part.usage.outputTokens
      }
    })
  }
}

// =============================================================================
// Tool Calling
// =============================================================================

/**
 * Represents all possible Anthropic provider-defined tools.
 *
 * @since 1.0.0
 * @category Models
 */
export type AnthropicTools =
  | typeof Generated.BetaTool.Encoded
  | typeof Generated.BetaBashTool20241022.Encoded
  | typeof Generated.BetaBashTool20250124.Encoded
  | typeof Generated.BetaComputerUseTool20241022.Encoded
  | typeof Generated.BetaComputerUseTool20250124.Encoded
  | typeof Generated.BetaTextEditor20241022.Encoded
  | typeof Generated.BetaTextEditor20250124.Encoded
  | typeof Generated.BetaTextEditor20250429.Encoded
  | typeof Generated.BetaTextEditor20250728.Encoded

/**
 * A helper method which takes in large language model provider options from
 * the base Effect AI SDK as well as Anthropic request configuration options
 * and returns the prepared tools, tool choice, and Anthropic betas to include
 * in a request.
 *
 * This method is primarily exposed for use by other Effect provider
 * integrations which can utilize Anthropic models (i.e. Amazon Bedrock).
 *
 * @since 1.0.0
 * @category Tool Calling
 */
export const prepareTools: (options: LanguageModel.ProviderOptions, config: Config.Service) => Effect.Effect<{
  readonly betas: ReadonlySet<string>
  readonly tools: ReadonlyArray<AnthropicTools> | undefined
  readonly toolChoice: typeof Generated.BetaToolChoice.Encoded | undefined
}, AiError.AiError> = Effect.fnUntraced(function*(options, config) {
  // Return immediately if no tools are in the toolkit or a tool choice of
  // "none" was specified
  if (options.tools.length === 0 || options.toolChoice === "none") {
    return { betas: new Set(), tools: undefined, toolChoice: undefined }
  }

  const betas = new Set<string>()
  let tools: Array<AnthropicTools> = []
  let toolChoice: typeof Generated.BetaToolChoice.Encoded | undefined = undefined

  // Convert the tools in the toolkit to the provider-defined format
  for (const tool of options.tools) {
    if (Tool.isUserDefined(tool)) {
      tools.push({
        name: tool.name,
        description: Tool.getDescription(tool as any),
        input_schema: Tool.getJsonSchema(tool as any) as any
      })
    }

    if (Tool.isProviderDefined(tool)) {
      switch (tool.id) {
        case "anthropic.bash_20241022": {
          betas.add("computer-use-2024-10-22")
          tools.push({
            name: "bash",
            type: "bash_20241022"
          })
          break
        }
        case "anthropic.bash_20250124": {
          betas.add("computer-use-2025-01-24")
          tools.push({
            name: "bash",
            type: "bash_20250124"
          })
          break
        }
        case "anthropic.code_execution_20250522": {
          betas.add("code-execution-2025-05-22")
          tools.push({
            ...tool.args,
            name: "code_execution",
            type: "code_execution_2025522"
          })
          break
        }
        case "anthropic.code_execution_20250825": {
          betas.add("code-execution-2025-08-25")
          tools.push({
            ...tool.args,
            name: "code_execution",
            type: "code_execution_20250825"
          })
          break
        }
        case "anthropic.computer_use_20241022": {
          betas.add("computer-use-2025-10-22")
          tools.push({
            ...tool.args,
            name: "computer",
            type: "computer_20241022"
          })
          break
        }
        case "anthropic.computer_use_20250124": {
          betas.add("computer-use-2025-01-24")
          tools.push({
            ...tool.args,
            name: "computer",
            type: "computer_20250124"
          })
          break
        }
        case "anthropic.text_editor_20241022": {
          betas.add("computer-use-2024-10-22")
          tools.push({
            name: "str_replace_editor",
            type: "text_editor_20241022"
          })
          break
        }
        case "anthropic.text_editor_20250124": {
          betas.add("computer-use-2025-01-24")
          tools.push({
            name: "str_replace_editor",
            type: "text_editor_20250124"
          })
          break
        }
        case "anthropic.text_editor_20250429": {
          betas.add("computer-use-2025-01-24")
          tools.push({
            name: "str_replace_based_edit_tool",
            type: "text_editor_20250429"
          })
          break
        }
        case "anthropic.text_editor_20250728": {
          tools.push({
            name: "str_replace_based_edit_tool",
            type: "text_editor_20250728"
          })
          break
        }
        case "anthropic.web_search_20250305": {
          tools.push({
            ...tool.args,
            name: "web_search",
            type: "web_search_20250305"
          })
          break
        }
        default: {
          return yield* new AiError.MalformedInput({
            module: "AnthropicLanguageModel",
            method: "prepareTools",
            description: `Received request to call unknown provider-defined tool '${tool.name}'`
          })
        }
      }
    }
  }

  // Convert the tool choice to the provider-defined format
  if (options.toolChoice === "auto") {
    toolChoice = {
      type: "auto",
      disable_parallel_tool_use: config.disableParallelToolCalls
    }
  } else if (options.toolChoice === "required") {
    toolChoice = {
      type: "any",
      disable_parallel_tool_use: config.disableParallelToolCalls
    }
  } else if ("tool" in options.toolChoice) {
    toolChoice = {
      type: "tool",
      name: options.toolChoice.tool,
      disable_parallel_tool_use: config.disableParallelToolCalls
    }
  } else {
    const allowedTools = new Set(options.toolChoice.oneOf)
    tools = tools.filter((tool) => allowedTools.has(tool.name))
    toolChoice = {
      type: options.toolChoice.mode === "required" ? "any" : "auto",
      disable_parallel_tool_use: config.disableParallelToolCalls
    }
  }

  return { betas, tools, toolChoice }
})

// =============================================================================
// Utilities
// =============================================================================

type ContentGroup = SystemMessageGroup | AssistantMessageGroup | UserMessageGroup

interface SystemMessageGroup {
  readonly type: "system"
  readonly messages: Array<Prompt.SystemMessage>
}

interface AssistantMessageGroup {
  readonly type: "assistant"
  readonly messages: Array<Prompt.AssistantMessage>
}

interface UserMessageGroup {
  readonly type: "user"
  readonly messages: Array<Prompt.ToolMessage | Prompt.UserMessage>
}

const groupMessages = (prompt: Prompt.Prompt): Array<ContentGroup> => {
  const messages: Array<ContentGroup> = []
  let current: ContentGroup | undefined = undefined
  for (const message of prompt.content) {
    switch (message.role) {
      case "system": {
        if (current?.type !== "system") {
          current = { type: "system", messages: [] }
          messages.push(current)
        }
        current.messages.push(message)
        break
      }
      case "assistant": {
        if (current?.type !== "assistant") {
          current = { type: "assistant", messages: [] }
          messages.push(current)
        }
        current.messages.push(message)
        break
      }
      case "tool":
      case "user": {
        if (current?.type !== "user") {
          current = { type: "user", messages: [] }
          messages.push(current)
        }
        current.messages.push(message)
        break
      }
    }
  }
  return messages
}

const isCitationPart = (part: Prompt.UserMessage["content"][number]): part is Prompt.FilePart => {
  if (part.type === "file" && (part.mediaType === "application/pdf" || part.mediaType === "text/plain")) {
    return part.options.anthropic?.citations?.enabled ?? false
  }
  return false
}

interface CitableDocument {
  readonly title: string
  readonly fileName: string | undefined
  readonly mediaType: string
}

const extractCitableDocuments = (prompt: Prompt.Prompt): ReadonlyArray<CitableDocument> => {
  const citableDocuments: Array<CitableDocument> = []
  for (const message of prompt.content) {
    if (message.role === "user") {
      for (const part of message.content) {
        if (isCitationPart(part)) {
          citableDocuments.push({
            title: part.fileName ?? "Untitled Document",
            fileName: part.fileName,
            mediaType: part.mediaType
          })
        }
      }
    }
  }
  return citableDocuments
}

const getCacheControl = (
  part:
    | Prompt.SystemMessage
    | Prompt.UserMessage
    | Prompt.AssistantMessage
    | Prompt.ToolMessage
    | Prompt.UserMessagePart
    | Prompt.AssistantMessagePart
    | Prompt.ToolMessagePart
): typeof Generated.CacheControlEphemeral.Encoded | undefined => part.options.anthropic?.cacheControl

const getDocumentMetadata = (part: Prompt.FilePart): {
  readonly title: string | undefined
  readonly context: string | undefined
} | undefined => {
  const options = part.options.anthropic
  if (Predicate.isNotUndefined(options)) {
    return {
      title: options.documentTitle,
      context: options.documentContext
    }
  }
  return undefined
}

const shouldEnableCitations = (part: Prompt.FilePart): boolean => part.options.anthropic?.citations?.enabled ?? false

const processCitation: (
  citation:
    | Generated.ResponseCharLocationCitation
    | Generated.ResponsePageLocationCitation
    | Generated.ResponseContentBlockLocationCitation
    | Generated.ResponseWebSearchResultLocationCitation
    | Generated.ResponseSearchResultLocationCitation,
  citableDocuments: ReadonlyArray<CitableDocument>,
  idGenerator: IdGenerator.Service
) => Effect.Effect<Response.DocumentSourcePartEncoded | Response.UrlSourcePartEncoded | undefined> = Effect.fnUntraced(
  function*(citation, citableDocuments, idGenerator) {
    if (citation.type === "page_location" || citation.type === "char_location") {
      const citedDocument = citableDocuments[citation.document_index]
      if (Predicate.isNotUndefined(citedDocument)) {
        const id = yield* idGenerator.generateId()

        const metadata = citation.type === "char_location"
          ? {
            source: "document",
            type: citation.type,
            citedText: citation.cited_text,
            startCharIndex: citation.start_char_index,
            endCharIndex: citation.end_char_index
          } as const
          : {
            source: "document",
            type: citation.type,
            citedText: citation.cited_text,
            startPageNumber: citation.start_page_number,
            endPageNumber: citation.end_page_number
          } as const

        return {
          type: "source",
          sourceType: "document",
          id,
          mediaType: citedDocument.mediaType,
          title: citation.document_title ?? citedDocument.title,
          fileName: citedDocument.fileName,
          metadata: { anthropic: metadata }
        }
      }
    }

    if (citation.type === "web_search_result_location") {
      const id = yield* idGenerator.generateId()

      const metadata = {
        source: "url",
        citedText: citation.cited_text,
        encryptedIndex: citation.encrypted_index
      } as const

      return {
        type: "source",
        sourceType: "url",
        id,
        url: citation.url,
        title: citation.title ?? "Untitled",
        metadata: { anthropic: metadata }
      }
    }
  }
)
